import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import requests
import datetime
from tqdm import tqdm

class ETFBacktest:
    def __init__(self, stock_code='510310.SH', start_date='20230101', end_date='20251231', 
                 initial_capital=500000, token="YOUR_ZTAPI_TOKEN"):
        """
        初始化回测系统
        
        Parameters:
        -----------
        stock_code: 股票代码，如 '510310.SH'
        start_date: 回测开始日期，格式 'YYYYMMDD'
        end_date: 回测结束日期，格式 'YYYYMMDD'
        initial_capital: 初始资金
        token: ZTAPI 的 token
        """
        self.stock_code = stock_code
        self.start_date = start_date
        self.end_date = end_date
        self.initial_capital = initial_capital
        self.token = token
        self.data = None
        self.positions = []  # 持仓记录 [数量, 买入价格]
        self.capital = initial_capital  # 可用资金
        self.transactions = []  # 交易记录
        self.pyramid_positions = []  # 金字塔网格补仓记录 [数量, 买入价格, 类型]
        self.last_pyramid_price = None  # 最后一次金字塔补仓价格
        self.history = []  # 资产历史
    
    def load_data_from_api(self):
        """
        从ZTAPI加载历史数据
        """
        print("从API加载数据...")
        url = f"https://api.zhituapi.com/hs/history/{self.stock_code}/d/n?token={self.token}&st={self.start_date}&et={self.end_date}"
        print("请求URL: "+url)
        response = requests.get(url)
        
        if response.status_code != 200:
            print(f"API请求失败，状态码: {response.status_code}")
            return False
        
        data = response.json()
        
        if not data:
            print("没有获取到数据")
            return False
        
        # 转换为DataFrame
        self.data = pd.DataFrame(data)
        self.data.rename(columns={
            't': 'date', 'o': 'open', 'h': 'high', 
            'l': 'low', 'c': 'close', 'v': 'volume', 
            'a': 'amount', 'pc': 'pre_close'
        }, inplace=True)
        
        # 转换日期格式
        self.data['date'] = pd.to_datetime(self.data['date'])
        
        print(f"成功加载 {len(self.data)} 天的数据")
        return True
    
    def calculate_indicators(self):
        """
        计算技术指标: KDJ(9,3,3), MACD(6,13,5), BOLL(20,2)
        """
        print("计算技术指标...")
        # 确保数据已加载
        if self.data is None:
            print("未加载数据，请先加载数据")
            return False
        
        # 计算KDJ
        low_list = self.data['low'].rolling(window=9).min()
        high_list = self.data['high'].rolling(window=9).max()
        
        rsv = (self.data['close'] - low_list) / (high_list - low_list) * 100
        self.data['K'] = rsv.ewm(alpha=1/3, adjust=False).mean()
        self.data['D'] = self.data['K'].ewm(alpha=1/3, adjust=False).mean()
        self.data['J'] = 3 * self.data['K'] - 2 * self.data['D']
        
        # 计算MACD
        self.data['ema6'] = self.data['close'].ewm(span=6, adjust=False).mean()
        self.data['ema13'] = self.data['close'].ewm(span=13, adjust=False).mean()
        self.data['diff'] = self.data['ema6'] - self.data['ema13']
        self.data['dea'] = self.data['diff'].ewm(span=5, adjust=False).mean()
        self.data['macd'] = 2 * (self.data['diff'] - self.data['dea'])
        
        # 计算BOLL
        self.data['ma20'] = self.data['close'].rolling(window=20).mean()
        self.data['std20'] = self.data['close'].rolling(window=20).std()
        self.data['upper'] = self.data['ma20'] + 2 * self.data['std20']
        self.data['lower'] = self.data['ma20'] - 2 * self.data['std20']
        
        # 计算成交量均值
        self.data['volume_ma20'] = self.data['volume'].rolling(window=20).mean()
        
        # 计算金叉死叉信号
        self.data['kdj_golden_cross'] = ((self.data['K'].shift(1) < self.data['D'].shift(1)) & 
                                          (self.data['K'] > self.data['D']))
        self.data['kdj_death_cross'] = ((self.data['K'].shift(1) > self.data['D'].shift(1)) & 
                                         (self.data['K'] < self.data['D']))
        
        self.data['macd_golden_cross'] = ((self.data['diff'].shift(1) < self.data['dea'].shift(1)) & 
                                           (self.data['diff'] > self.data['dea']))
        self.data['macd_death_cross'] = ((self.data['diff'].shift(1) > self.data['dea'].shift(1)) & 
                                          (self.data['diff'] < self.data['dea']))
        
        # 丢弃NaN值
        self.data = self.data.dropna()
        return True
    
    def check_buy_signal(self, row):
        """
        检查买入信号（放宽条件）
        - MACD在0轴下方
        - J值小于20（原来是10，放宽了）
        - 出现MACD或KDJ金叉
        
        返回:
        - True: 满足买入条件
        - False: 不满足买入条件
        """
        # MACD在0轴下方
        macd_below_zero = row['diff'] < 0
        
        # J值小于20（放宽条件）
        j_below_20 = row['J'] < 20
        
        # 出现MACD或KDJ金叉
        golden_cross = row['macd_golden_cross'] or row['kdj_golden_cross']
        
        # 检查是否符合所有条件
        buy_signal = macd_below_zero and j_below_20 and golden_cross
        
        # 如果满足买入条件，记录详细信息
        if buy_signal:
            print(f"【买入信号】 {row['date']}: 价格={row['close']}, J={row['J']:.2f}, MACD差值={row['diff']:.5f}")
        
        return buy_signal
    
    def check_sell_signal(self, row, prev_row):
        """
        检查卖出信号
        卖出信号: MACD0轴上方，且J值大于80分5批出售；如果前一日J值大于80，前一日卖出1批了，当但日小于80，则当日不卖
        """
        # MACD在0轴上方
        macd_above_zero = row['diff'] > 0
        
        # J值大于80
        j_above_80 = row['J'] > 80
        
        # 前一日J值大于80
        prev_j_above_80 = prev_row['J'] > 80 if prev_row is not None else False
        
        # 当日J值小于80
        j_below_80 = row['J'] <= 80
        
        # 如果前一日J值大于80，前一日已卖出，当日J值小于80，则当日不卖
        if prev_j_above_80 and j_below_80:
            return 0  # 不卖出
        
        # MACD在0轴上方且J值大于80，分5批卖出
        if macd_above_zero and j_above_80:
            return 1  # 卖出1批
        
        return 0  # 不卖出
    
    def execute_buy(self, price, date):
        """
        执行买入操作
        买入数量为10000股
        """
        buy_shares = 10000
        cost = buy_shares * price
        
        if self.capital >= cost:
            self.capital -= cost
            self.positions.append([buy_shares, price])
            self.transactions.append({
                'date': date,
                'type': 'buy',
                'price': price,
                'shares': buy_shares,
                'cost': cost,
                'remaining_capital': self.capital
            })
            print(f"{date} 买入: {buy_shares}股, 价格: {price}, 花费: {cost:.2f}, 剩余资金: {self.capital:.2f}")
            return True
        else:
            print(f"{date} 资金不足，无法买入")
            return False
    
    def execute_sell(self, price, date, batch_size=0.2):
        """
        执行卖出操作
        分5批卖出（每批20%）
        """
        if not self.positions:
            return False
        
        # 计算要卖出的股数
        total_shares = sum(pos[0] for pos in self.positions)
        sell_shares = int(total_shares * batch_size)
        
        if sell_shares <= 0:
            return False
        
        # 计算收入
        income = sell_shares * price
        
        # 按照FIFO原则卖出股票
        shares_to_sell = sell_shares
        avg_buy_price = 0
        shares_removed = 0
        
        while shares_to_sell > 0 and self.positions:
            pos_shares, pos_price = self.positions[0]
            
            if pos_shares <= shares_to_sell:
                # 全部卖出这个持仓
                self.positions.pop(0)
                shares_to_sell -= pos_shares
                avg_buy_price += pos_price * pos_shares
                shares_removed += pos_shares
            else:
                # 部分卖出这个持仓
                self.positions[0][0] -= shares_to_sell
                avg_buy_price += pos_price * shares_to_sell
                shares_removed += shares_to_sell
                shares_to_sell = 0
        
        avg_buy_price = avg_buy_price / shares_removed if shares_removed > 0 else 0
        profit = income - avg_buy_price * sell_shares
        
        self.capital += income
        self.transactions.append({
            'date': date,
            'type': 'sell',
            'price': price,
            'shares': sell_shares,
            'income': income,
            'profit': profit,
            'remaining_capital': self.capital
        })
        
        print(f"{date} 卖出: {sell_shares}股, 价格: {price}, 收入: {income:.2f}, 利润: {profit:.2f}, 剩余资金: {self.capital:.2f}")
        return True
    
    def pyramid_strategy(self, current_price, date):
        """
        金字塔网格补仓策略
        - 初次买入后，下跌3%，则补仓初次买入数量的1.2倍；后续每下跌3%，都补上一次补仓数量的1.2倍
        - 当出现反转时，从最后一次补仓上涨3%开始，卖出补仓数量，第一次卖出数量为10000股，后续每上涨3%卖出为上次补仓数量的1.2倍
        """
        # 如果从未进行过交易，则不执行金字塔策略
        if not self.pyramid_positions and not self.positions:
            return
        
        # 确定最后一次操作的价格
        if self.last_pyramid_price is None:
            if self.positions:
                # 使用首次买入的价格
                self.last_pyramid_price = self.positions[0][1]
            else:
                return
        
        # 计算价格变动百分比
        price_change = (current_price - self.last_pyramid_price) / self.last_pyramid_price
        
        # 下跌3%则补仓
        if price_change <= -0.03:
            # 计算补仓数量
            if not self.pyramid_positions:
                # 首次补仓，为初次买入数量的1.2倍
                buy_shares = int(10000 * 1.2)
            else:
                # 后续补仓，为上次补仓数量的1.2倍
                last_shares = self.pyramid_positions[-1][0]
                buy_shares = int(last_shares * 1.2)
            
            cost = buy_shares * current_price
            
            if self.capital >= cost:
                self.capital -= cost
                self.pyramid_positions.append([buy_shares, current_price, 'buy'])
                self.last_pyramid_price = current_price
                self.transactions.append({
                    'date': date,
                    'type': 'pyramid_buy',
                    'price': current_price,
                    'shares': buy_shares,
                    'cost': cost,
                    'remaining_capital': self.capital
                })
                print(f"{date} 金字塔补仓: {buy_shares}股, 价格: {current_price}, 花费: {cost:.2f}, 剩余资金: {self.capital:.2f}")
            else:
                print(f"{date} 资金不足，无法进行金字塔补仓")
        
        # 上涨3%则卖出
        elif price_change >= 0.03:
            if self.pyramid_positions:
                # 找到最近一次补仓
                last_pyramid = None
                for pos in reversed(self.pyramid_positions):
                    if pos[2] == 'buy':
                        last_pyramid = pos
                        break
                
                if last_pyramid:
                    # 计算卖出数量
                    sell_shares = int(last_pyramid[0] * 1.2)
                    
                    # 检查是否有足够的股票可以卖出
                    total_shares = sum(pos[0] for pos in self.positions)
                    total_shares += sum(pos[0] for pos in self.pyramid_positions if pos[2] == 'buy')
                    total_shares -= sum(pos[0] for pos in self.pyramid_positions if pos[2] == 'sell')
                    
                    if total_shares >= sell_shares:
                        income = sell_shares * current_price
                        self.capital += income
                        self.pyramid_positions.append([sell_shares, current_price, 'sell'])
                        self.last_pyramid_price = current_price
                        self.transactions.append({
                            'date': date,
                            'type': 'pyramid_sell',
                            'price': current_price,
                            'shares': sell_shares,
                            'income': income,
                            'remaining_capital': self.capital
                        })
                        print(f"{date} 金字塔卖出: {sell_shares}股, 价格: {current_price}, 收入: {income:.2f}, 剩余资金: {self.capital:.2f}")
                    else:
                        print(f"{date} 持仓不足，无法进行金字塔卖出")
                else:
                    # 如果没有找到补仓记录但有初始买入，则卖出10000股
                    if self.positions:
                        sell_shares = 10000
                        total_shares = sum(pos[0] for pos in self.positions)
                        
                        if total_shares >= sell_shares:
                            income = sell_shares * current_price
                            self.capital += income
                            self.pyramid_positions.append([sell_shares, current_price, 'sell'])
                            self.last_pyramid_price = current_price
                            self.transactions.append({
                                'date': date,
                                'type': 'pyramid_sell',
                                'price': current_price,
                                'shares': sell_shares,
                                'income': income,
                                'remaining_capital': self.capital
                            })
                            print(f"{date} 金字塔初次卖出: {sell_shares}股, 价格: {current_price}, 收入: {income:.2f}, 剩余资金: {self.capital:.2f}")
                        else:
                            print(f"{date} 持仓不足，无法进行金字塔卖出")
    
    def calculate_total_assets(self, price):
        """
        计算总资产价值
        """
        # 计算现有持仓市值
        position_value = sum(pos[0] * price for pos in self.positions)
        
        # 计算金字塔持仓市值
        pyramid_value = 0
        if self.pyramid_positions:
            buy_shares = sum(pos[0] for pos in self.pyramid_positions if pos[2] == 'buy')
            sell_shares = sum(pos[0] for pos in self.pyramid_positions if pos[2] == 'sell')
            pyramid_value = (buy_shares - sell_shares) * price
        
        # 总资产 = 可用资金 + 持仓市值 + 金字塔持仓市值
        total_assets = self.capital + position_value + pyramid_value
        return total_assets
    
    def run_backtest(self):
        """
        运行回测
        """
        print("开始回测...")
        if self.data is None or len(self.data) == 0:
            print("没有数据可用于回测")
            return False
        
        # 确保已计算指标
        if 'J' not in self.data.columns:
            self.calculate_indicators()
        
        # 打印数据统计信息，帮助诊断问题
        print("\n===== 数据统计信息 =====")
        print(f"总数据点数: {len(self.data)}")
        print(f"价格范围: {self.data['close'].min()} - {self.data['close'].max()}")
        print(f"J值范围: {self.data['J'].min()} - {self.data['J'].max()}")
        print(f"MACD diff范围: {self.data['diff'].min()} - {self.data['diff'].max()}")
        
        # 检查买入条件满足的天数
        macd_below_zero_count = sum(self.data['diff'] < 0)
        j_below_10_count = sum(self.data['J'] < 10)
        golden_cross_count = sum(self.data['macd_golden_cross'] | self.data['kdj_golden_cross'])
        volume_breakthrough_count = sum(self.data['volume'] > 1.5 * self.data['volume_ma20'])
        
        print(f"MACD在0轴下方的天数: {macd_below_zero_count}")
        print(f"J值小于10的天数: {j_below_10_count}")
        print(f"出现金叉的天数: {golden_cross_count}")
        print(f"成交量突破的天数: {volume_breakthrough_count}")
        print("=========================\n")
        
        # 初始化
        self.positions = []
        self.capital = self.initial_capital
        self.transactions = []  # 清空交易记录列表
        self.pyramid_positions = []
        self.last_pyramid_price = None
        self.history = []
        
        # 前一天的数据
        prev_row = None
        
        # 添加计数器，用于跟踪符合信号的日期数
        buy_signals_count = 0
        sell_signals_count = 0
        
        # 模拟每天的交易
        for i, row in tqdm(self.data.iterrows(), total=len(self.data)):
            date = row['date']
            current_price = row['close']
            
            # 检查买入信号
            if self.check_buy_signal(row):
                buy_signals_count += 1
                print(f"【买入信号】 {date}: 价格={current_price}, J={row['J']:.2f}, MACD差值={row['diff']:.5f}")
                self.execute_buy(current_price, date)
            
            # 检查卖出信号
            sell_batch = self.check_sell_signal(row, prev_row)
            if sell_batch > 0:
                sell_signals_count += 1
                print(f"【卖出信号】 {date}: 价格={current_price}, J={row['J']:.2f}, MACD差值={row['diff']:.5f}")
                self.execute_sell(current_price, date)
            
            # 执行金字塔网格补仓策略
            self.pyramid_strategy(current_price, date)
            
            # 保存每日资产状况
            total_assets = self.calculate_total_assets(current_price)
            self.history.append({
                'date': date,
                'price': current_price,
                'capital': self.capital,
                'assets': total_assets
            })
            
            # 更新前一天的数据
            prev_row = row
        
        # 统计结果
        print(f"\n总共产生买入信号: {buy_signals_count} 次")
        print(f"总共产生卖出信号: {sell_signals_count} 次")
        
        # 转换历史记录为DataFrame
        self.history = pd.DataFrame(self.history)
        
        # 只有当有交易记录时才转换成DataFrame
        if self.transactions:
            self.transactions = pd.DataFrame(self.transactions)
        else:
            # 创建一个空的DataFrame但包含必要的列
            self.transactions = pd.DataFrame(columns=['date', 'type', 'price', 'shares', 'cost', 'income', 'profit', 'remaining_capital'])
    
        print("回测完成!")
        return True
    
    def calculate_metrics(self):
        """
        计算绩效指标
        """
        if len(self.history) == 0:
            print("没有历史数据，无法计算绩效")
            return {}
        
        start_assets = self.initial_capital
        end_assets = self.history.iloc[-1]['assets']
        
        # 计算总收益
        total_return = (end_assets - start_assets) / start_assets
        
        # 计算年化收益率
        days = (self.history.iloc[-1]['date'] - self.history.iloc[0]['date']).days
        years = days / 365
        annual_return = (1 + total_return) ** (1 / years) - 1 if years > 0 else 0
        
        # 计算最大回撤
        self.history['drawdown'] = (self.history['assets'].cummax() - self.history['assets']) / self.history['assets'].cummax()
        max_drawdown = self.history['drawdown'].max()
        
        # 计算夏普比率 (假设无风险利率为3%)
        daily_returns = self.history['assets'].pct_change().dropna()
        avg_return = daily_returns.mean()
        risk_free_rate = 0.03 / 252  # 日度无风险利率
        sharpe_ratio = (avg_return - risk_free_rate) / daily_returns.std() * (252 ** 0.5) if daily_returns.std() > 0 else 0
        
        # 计算交易次数
        buy_trades = sell_trades = pyramid_buy_trades = pyramid_sell_trades = 0
        
        if hasattr(self, 'transactions') and isinstance(self.transactions, pd.DataFrame) and not self.transactions.empty:
            # 检查列是否存在
            if 'type' in self.transactions.columns:
                buy_trades = len(self.transactions[self.transactions['type'] == 'buy'])
                sell_trades = len(self.transactions[self.transactions['type'] == 'sell'])
                pyramid_buy_trades = len(self.transactions[self.transactions['type'] == 'pyramid_buy'])
                pyramid_sell_trades = len(self.transactions[self.transactions['type'] == 'pyramid_sell'])
            else:
                print("警告: 交易记录中没有'type'列，无法计算交易次数")
        
        metrics = {
            'start_date': self.history.iloc[0]['date'],
            'end_date': self.history.iloc[-1]['date'],
            'start_assets': start_assets,
            'end_assets': end_assets,
            'total_return': total_return,
            'annual_return': annual_return,
            'max_drawdown': max_drawdown,
            'sharpe_ratio': sharpe_ratio,
            'buy_trades': buy_trades,
            'sell_trades': sell_trades,
            'pyramid_buy_trades': pyramid_buy_trades,
            'pyramid_sell_trades': pyramid_sell_trades,
            'total_trades': buy_trades + sell_trades + pyramid_buy_trades + pyramid_sell_trades
        }
        
        return metrics
    
    def plot_results(self):
        """
        绘制回测结果
        """
        if len(self.history) == 0:
            print("没有历史数据，无法绘图")
            return
        
        plt.figure(figsize=(14, 10))
        
        # 创建上下两个子图，比例为3:1
        gs = plt.GridSpec(4, 1)
        price_ax = plt.subplot(gs[0:3, 0])
        assets_ax = plt.subplot(gs[3, 0], sharex=price_ax)
        
        # 绘制股票价格走势
        price_ax.plot(self.history['date'], self.history['price'], label='Close Price')
        price_ax.set_ylabel('Price')
        price_ax.set_title(f'{self.stock_code} Backtest Results')
        price_ax.legend(loc='upper left')
        
        # 在价格图上标记买卖点
        if hasattr(self, 'transactions') and isinstance(self.transactions, pd.DataFrame):
            buys = self.transactions[self.transactions['type'] == 'buy']
            sells = self.transactions[self.transactions['type'] == 'sell']
            pyramid_buys = self.transactions[self.transactions['type'] == 'pyramid_buy']
            pyramid_sells = self.transactions[self.transactions['type'] == 'pyramid_sell']
            
            for _, trade in buys.iterrows():
                price_ax.plot(trade['date'], trade['price'], '^', markersize=8, color='green')
            
            for _, trade in sells.iterrows():
                price_ax.plot(trade['date'], trade['price'], 'v', markersize=8, color='red')
            
            for _, trade in pyramid_buys.iterrows():
                price_ax.plot(trade['date'], trade['price'], 's', markersize=6, color='blue')
            
            for _, trade in pyramid_sells.iterrows():
                price_ax.plot(trade['date'], trade['price'], 'd', markersize=6, color='purple')
        
        # 绘制资产走势
        assets_ax.plot(self.history['date'], self.history['assets'], label='Total Assets')
        assets_ax.set_xlabel('Date')
        assets_ax.set_ylabel('Assets')
        assets_ax.legend(loc='upper left')
        
        plt.tight_layout()
        plt.show()
        
        # 绘制回撤图
        plt.figure(figsize=(14, 6))
        plt.plot(self.history['date'], self.history['drawdown'] * 100)
        plt.fill_between(self.history['date'], 0, self.history['drawdown'] * 100, alpha=0.3, color='red')
        plt.xlabel('Date')
        plt.ylabel('Drawdown (%)')
        plt.title('Portfolio Drawdown')
        plt.tight_layout()
        plt.show()
    
    def print_summary(self):
        """
        打印回测摘要
        """
        metrics = self.calculate_metrics()
        
        print("\n========== 回测摘要 ==========")
        print(f"股票代码: {self.stock_code}")
        print(f"回测周期: {metrics['start_date']} 至 {metrics['end_date']}")
        print(f"初始资金: {metrics['start_assets']:.2f}")
        print(f"最终资产: {metrics['end_assets']:.2f}")
        print(f"总收益率: {metrics['total_return']*100:.2f}%")
        print(f"年化收益率: {metrics['annual_return']*100:.2f}%")
        print(f"最大回撤: {metrics['max_drawdown']*100:.2f}%")
        print(f"夏普比率: {metrics['sharpe_ratio']:.2f}")
        print(f"买入交易次数: {metrics['buy_trades']}")
        print(f"卖出交易次数: {metrics['sell_trades']}")
        print(f"金字塔补仓次数: {metrics['pyramid_buy_trades']}")
        print(f"金字塔卖出次数: {metrics['pyramid_sell_trades']}")
        print(f"总交易次数: {metrics['total_trades']}")
        print("================================")


# 示例使用
if __name__ == "__main__":
    # 初始化回测系统
    backtest = ETFBacktest(
        stock_code='512880.SH',
        start_date='20230101',  
        end_date='20251231',
        initial_capital=500000,
        token="FF25A28E-2FE6-4047-9691-A8E2B9E58A61"  # 请替换成你的实际token
    )
    
    # 加载数据
    data_loaded = backtest.load_data_from_api()
    
    if data_loaded:
        # 计算技术指标
        backtest.calculate_indicators()
        
        # 运行回测
        backtest.run_backtest()
        
        # 打印结果摘要
        backtest.print_summary()
        
        # 绘制回测结果
        backtest.plot_results()
